#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import glob
import os.path
import sys
import time

sys.path.insert(0, './gen-py')
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
lib_path = glob.glob('../../lib/py/build/lib.*')
if lib_path:
    sys.path.insert(0, lib_path[0])
from optparse import OptionParser

from ThriftTest import ThriftTest, SecondService
from ThriftTest.ttypes import *
from thrift import TMultiplexedProcessor
from thrift.Thrift import TProcessorEventHandler
from thrift.transport import TTransport
from thrift.transport import TSocket, TSSLSocket
from thrift.transport.THeaderTransport import THeaderTransport
from thrift.protocol import TBinaryProtocol
from thrift.protocol import THeaderProtocol
from thrift.server import TServer, TNonblockingServer, \
    TProcessPoolServer, THttpServer, TCppServer

class SecondHandler(SecondService.Iface):
    def blahBlah(self):
        print('blahBlah()')

class SecondContextHandler(SecondService.ContextIface):
    def __init__(self,):
        self.th = SecondHandler()

    def blahBlah(self, handler_ctx):
        self.th.blahBlah()

class TestHandler(ThriftTest.Iface):

    def testVoid(self):
        print('testVoid()')

    def testString(self, str):
        print('testString(%s)' % str)
        return str

    def testByte(self, byte):
        print('testByte(%d)' % byte)
        return byte

    def testI16(self, i16):
        print('testI16(%d)' % i16)
        return i16

    def testI32(self, i32):
        print('testI32(%d)' % i32)
        return i32

    def testI64(self, i64):
        print('testI64(%d)' % i64)
        return i64

    def testDouble(self, dub):
        print('testDouble(%f)' % dub)
        return dub

    def testFloat(self, flt):
        print('testFloat(%f)' % flt)
        return flt

    def testStruct(self, thing):
        print('testStruct({%s, %d, %d, %d})' % (thing.string_thing,
            thing.byte_thing, thing.i32_thing, thing.i64_thing))
        return thing

    def testException(self, str):
        print('testException(%s)' % str)
        if str == 'Xception':
            x = Xception()
            x.errorCode = 1001
            x.message = str
            raise x
        elif str == "throw_undeclared":
            raise ValueError("foo")

    def testOneway(self, seconds):
        print('testOneway(%d) => sleeping...' % seconds)
        time.sleep(seconds)
        print('done sleeping')

    def testNest(self, thing):
        return thing

    def testMap(self, thing):
        return thing

    def testSet(self, thing):
        return thing

    def testList(self, thing):
        return thing

    def testEnum(self, thing):
        return thing

    def testTypedef(self, thing):
        return thing

    def testRequestCount(self):
        return event_handler.request_count

    def testNewConnection(self):
        return event_handler.num_new_conns

    def testPreServe(self):
        return event_handler.num_pre_serve

    def testConnectionDestroyed(self):
        return event_handler.num_conns_destroyed


class TestContextHandler(ThriftTest.ContextIface):

    def __init__(self, server_port):
        self.th = TestHandler()
        self._server_port = server_port

    def testVoid(self, handler_ctx):
        self.th.testVoid()
        # This is here so we can check that handler_ctx is getting set,
        # without modifying the service definition which would require
        # modifying all the languages.
        if not (handler_ctx[0].endswith("127.0.0.1") or \
            handler_ctx[0].endswith("::1")) or \
                handler_ctx[1] == self._server_port:
            raise ValueError("handler_ctx not set properly " +
                    str(handler_ctx))

    def testString(self, handler_ctx, str):
        return self.th.testString(str)

    def testByte(self, handler_ctx, byte):
        return self.th.testByte(byte)

    def testI16(self, handler_ctx, i16):
        return self.th.testI16(i16)

    def testI32(self, handler_ctx, i32):
        return self.th.testI32(i32)

    def testI64(self, handler_ctx, i64):
        return self.th.testI64(i64)

    def testDouble(self, handler_ctx, dub):
        return self.th.testDouble(dub)

    def testFloat(self, handler_ctx, flt):
        return self.th.testFloat(flt)

    def testStruct(self, handler_ctx, thing):
        return self.th.testStruct(thing)

    def testException(self, handler_ctx, str):
        return self.th.testException(str)

    def testOneway(self, handler_ctx, seconds):
        return self.th.testOneway(seconds)

    def testNest(self, handler_ctx, thing):
        return self.th.testNest(thing)

    def testMap(self, handler_ctx, thing):
        return self.th.testMap(thing)

    def testSet(self, handler_ctx, thing):
        return self.th.testSet(thing)

    def testList(self, handler_ctx, thing):
        return self.th.testList(thing)

    def testEnum(self, handler_ctx, thing):
        return self.th.testEnum(thing)

    def testTypedef(self, handler_ctx, thing):
        return self.th.testTypedef(thing)

    def testRequestCount(self, handler_ctx):
        return self.th.testRequestCount()

    def testNewConnection(self, handler_ctx):
        return self.th.testNewConnection()

    def testPreServe(self, handler_ctx):
        return self.th.testPreServe()

    def testConnectionDestroyed(self, handler_ctx):
        return self.th.testConnectionDestroyed()


class ContextEventHandler(TProcessorEventHandler):
    def getHandlerContext(self, fn_name, server_context):
        # this is a tuple ("hostname", port)
        return server_context.getPeerName()

class HeaderEventHandler(ContextEventHandler):
    def getHandlerContext(self, fn_name, server_context):
        self.htrans = server_context.iprot.trans
        return ContextEventHandler.getHandlerContext(self,
                fn_name, server_context)

    def preWrite(self, handler_context, fn_name, result):
        for str_key, str_value in self.htrans.get_headers().items():
            # Just spit them back for testing.
            self.htrans.set_header(str_key, str_value)

class TestServerEventHandler(TServer.TServerEventHandler):

    def __init__(self):
        self.num_pre_serve = 0
        self.request_count = 0
        self.num_new_conns = 0
        self.num_conns_destroyed = 0

    def newConnection(self, context):
        self.num_new_conns += 1

    def preServe(self, address):
        self.num_pre_serve += 1

    def clientBegin(self, iprot, oprot):
        self.request_count += 1

    def connectionDestroyed(self, context):
        self.num_conns_destroyed += 1


if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option(
        "--ssl",
        action="store_true",
        dest="ssl",
        default=False,
        help="use SSL for encrypted transport")
    parser.add_option(
        "--multiple",
        action="store_true",
        dest="multiple",
        default=False,
        help="use multiple service")
    parser.add_option(
        "--header",
        action="store_true",
        dest="header",
        default=False,
        help="use the Header protocol")
    parser.add_option(
        "--context",
        action="store_true",
        dest="context",
        default=False,
        help="Use the context-passing Handler")
    parser.add_option(
        "--port",
        action="store",
        type="int",
        dest="port",
        default=9090)
    parser.add_option(
        "--timeout",
        action="store",
        type="int",
        dest="timeout",
        default=60)
    options, args = parser.parse_args()

    event_handler = TestServerEventHandler()

    if options.header:
        pfactory = THeaderProtocol.THeaderProtocolFactory(True,
                [THeaderTransport.HEADERS_CLIENT_TYPE,
                 THeaderTransport.FRAMED_DEPRECATED,
                 THeaderTransport.UNFRAMED_DEPRECATED,
                 THeaderTransport.HTTP_CLIENT_TYPE])
    else:
        pfactory = TBinaryProtocol.TBinaryProtocolFactory()

    if options.context:
        processor = ThriftTest.ContextProcessor(TestContextHandler(
                options.port))
    else:
        processor = ThriftTest.Processor(TestHandler())

    if options.multiple:
        processor = TMultiplexedProcessor.TMultiplexedProcessor()
        if options.context:
            processor.registerProcessor("ThriftTest",
                    ThriftTest.ContextProcessor(TestContextHandler(
                            options.port)))
            processor.registerProcessor("SecondService",
                    SecondService.ContextProcessor(SecondContextHandler()))
        else:
            processor.registerProcessor("ThriftTest",
                    ThriftTest.Processor(TestHandler()))
            processor.registerProcessor("SecondService",
                    SecondService.Processor(SecondHandler()))

    if args[0] == "THttpServer":
        try:
            server = THttpServer.THttpServer(
                    processor, ('', options.port), pfactory)
        except:
            print("Could not load THttpServer")
            raise
    elif args[0] == "TCppServer":
        server = TCppServer.TCppServer(processor)
        server.setPort(options.port)
    else:
        if options.ssl:
            cert_path = 'thrift/test/py/test_cert.pem'
            transport = TSSLSocket.TSSLServerSocket(options.port,
                    certfile=cert_path)
        else:
            transport = TSocket.TServerSocket(options.port)
        tfactory = TTransport.TBufferedTransportFactory()

        server_class_name = args[0]

        if args[0] == "TNonblockingServer":
            server = TNonblockingServer.TNonblockingServer(
                    processor, transport, pfactory,
                    readTimeout=options.timeout)
        elif args[0] == "TProcessPoolServer":
            server = TProcessPoolServer.TProcessPoolServer(
                    processor, transport, tfactory, pfactory)
        else:
            ServerClass = getattr(TServer, server_class_name)
            server = ServerClass(processor, transport, tfactory, pfactory)

    if options.header:
        server.processor.setEventHandler(HeaderEventHandler())
    elif options.context:
        server.processor.setEventHandler(ContextEventHandler())
    server.setServerEventHandler(event_handler)

    server.serve()
